from typing import List, Dict, Any, Optional, Union
from langchain.tools import Tool
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.agents import AgentExecutor
from langchain.agents.output_parsers import OpenAIFunctionsAgentOutputParser
from langchain.agents.format_scratchpad import format_to_openai_functions
from langchain_core.messages import SystemMessage
from langchain.agents import create_tool_calling_agent
from app.core.logger import get_logger
from app.agent.db_tool import DatabaseTool, mcp_list_databases, mcp_execute_query, mcp_get_table_names, mcp_get_table_schema
from app.agent.db_schema_reader import DatabaseSchemaReader
from app.agent.finance_tool import FinanceDatabaseTool
from app.agent.llm import get_llm

logger = get_logger(__name__)

class Agent:
    """智能体基类"""
    
    def __init__(self, name: str, description: str = ""):
        self.name = name
        self.description = description
        self.tools: List[Tool] = []
        self.llm = get_llm()
        self.agent_executor = None
        self.db_tool = DatabaseTool()
        
    def add_tool(self, tool: Tool):
        """添加工具到智能体"""
        self.tools.append(tool)
        
    def get_system_prompt(self) -> str:
        """获取系统提示词"""
        return """你是一个智能助手，可以使用各种工具来帮助用户解决问题。"""
    
    def initialize_agent(self):
        """初始化智能体"""
        raise NotImplementedError("子类必须实现initialize_agent方法")
    
    def execute(self, query: str, config: Dict[str, Any] = None) -> str:
        """执行查询，增加错误处理"""
        # 为每次查询动态创建agent_executor
        if not self.agent_executor:
            self.initialize_agent()
            
        if config is None:
            config = {"configurable": {"thread_id": "finance-agent-thread"}}
            
        try:
            result = self.agent_executor.invoke({"input": query}, config=config)
            # 处理不同格式的结果
            if isinstance(result, dict):
                # 尝试获取output字段
                if "output" in result:
                    return result["output"]
                # 尝试获取其他可能的字段
                elif "result" in result:
                    return result["result"]
                elif "response" in result:
                    return result["response"]
                # 如果都没有，返回整个结果的字符串表示
                return str(result)
            # 如果结果不是字典，直接返回
            return str(result)
        except Exception as e:
            logger.error(f"执行查询时出错: {e}")
            
            # 提供更友好的错误信息
            error_msg = f"抱歉，处理您的查询时出现错误: {str(e)}。"
            
            # 检查是否有finance_db_tool属性，如果有则尝试直接执行SQL查询
            if hasattr(self, 'finance_db_tool'):
                try:
                    # 简单的SQL提取（这只是一个示例，可以根据具体需求完善）
                    import re
                    sql_match = re.search(r'(SELECT\s+.*?FROM\s+.*?)(?:\s+WHERE\s+.*?);?', query, re.IGNORECASE | re.DOTALL)
                    
                    if sql_match:
                        sql_query = sql_match.group(1).strip()
                        # 使用execute_finance_query工具直接执行查询
                        result = self.finance_db_tool.execute_query(sql_query)
                        
                        if result:
                            return f"找到结果：{str(result)}"
                        else:
                            return "查询没有返回结果。请检查您的查询条件或尝试其他查询。"
                    else:
                        return error_msg + " 无法从您的查询中提取SQL。请确保您的查询格式正确。"
                        
                except Exception as inner_e:
                    logger.error(f"备选查询方法出错: {inner_e}")
                    return error_msg + " 尝试直接访问数据库也失败了。"
            
            return error_msg

class DatabaseAgent(Agent):
    """数据库智能体"""
    
    def __init__(self, name: str = "Database Agent", description: str = "可以访问和查询sqlite数据库的智能体"):
        super().__init__(name, description)
        self._init_tools()
        
    def _init_tools(self):
        """初始化数据库相关工具"""
        # 列出数据库工具
        list_db_tool = Tool.from_function(
            func=self.db_tool.list_databases,
            name="list_databases",
            description="列出所有可用的数据库文件"
        )
        
        # 执行查询工具
        def execute_query_tool(db_name: str, query: str) -> List[Dict[str, Any]]:
            """执行SQL查询"""
            return self.db_tool.execute_query(db_name, query)
        
        execute_query_tool_wrapper = Tool.from_function(
            func=execute_query_tool,
            name="execute_query",
            description="在指定数据库上执行SQL查询。参数：db_name(数据库名), query(SQL查询语句)"
        )
        
        # 获取表名工具
        def get_table_names_tool(db_name: str) -> List[str]:
            """获取数据库中所有表名"""
            return self.db_tool.get_table_names(db_name)
        
        get_table_names_tool_wrapper = Tool.from_function(
            func=get_table_names_tool,
            name="get_table_names",
            description="获取指定数据库中所有表名。参数：db_name(数据库名)"
        )
        
        # 获取表结构工具
        def get_table_schema_tool(db_name: str, table_name: str) -> List[Dict[str, Any]]:
            """获取表的结构信息"""
            return self.db_tool.get_table_schema(db_name, table_name)
        
        get_table_schema_tool_wrapper = Tool.from_function(
            func=get_table_schema_tool,
            name="get_table_schema",
            description="获取指定表的结构信息。参数：db_name(数据库名), table_name(表名)"
        )
        
        self.tools = [
            list_db_tool,
            execute_query_tool_wrapper,
            get_table_names_tool_wrapper,
            get_table_schema_tool_wrapper
        ]
        
    def get_system_prompt(self) -> str:
        """获取数据库智能体的系统提示词"""
        return """
你是一个经验丰富的数据库专家助手，可以高效访问和查询sqlite数据库。你可以使用以下工具：

1. list_databases - 列出所有可用的数据库文件，返回格式为文件名列表
2. get_table_names - 获取指定数据库中所有表名，参数：db_name(数据库名)
3. get_table_schema - 获取指定表的结构信息，包括列名、数据类型等，参数：db_name(数据库名), table_name(表名)
4. execute_query - 在指定数据库上执行SQL查询，参数：db_name(数据库名), query(SQL查询语句)

操作流程规范（请严格遵循）：
1. 对于任何查询，第一步必须使用list_databases工具获取所有可用的数据库文件
2. 第二步使用get_table_names工具查看目标数据库中的所有表名
3. 第三步使用get_table_schema工具查看与查询相关表的详细结构
4. 第四步根据获取的表结构信息，构建语法正确的SQL查询语句
5. 第五步使用execute_query工具执行查询并获取结果

SQL查询最佳实践：
1. 所有SQL查询必须是SELECT语句，绝对禁止执行INSERT、UPDATE、DELETE、DROP等修改数据的操作
2. 构建查询语句时，必须严格按照表结构中的列名，对于包含特殊字符的列名，使用双引号包裹
3. 查询结果最多返回10条记录，使用LIMIT子句限制结果数量
4. 使用ORDER BY子句对结果进行排序，使返回的数据更有价值
5. 对于复杂查询，使用WHERE子句精确过滤数据

回答标准：
1. 回答必须基于实际的查询结果，确保信息的准确性
2. 回答时要简明扼要，直接针对用户问题提供答案，避免冗余信息
3. 如果用户没有指定具体的数据库，必须先列出所有可用的数据库供用户选择
4. 如果查询没有返回结果，明确告知用户未找到相关数据
5. 如果查询过程中出现错误，详细说明错误原因并提供可行的解决方案
6. 不要在回答中直接显示SQL语句，只需要解释查询结果即可

现在开始处理用户的查询，严格按照上述流程和规范执行。"""

    def initialize_agent(self):
        """初始化数据库智能体"""
        try:
            # 尝试使用create_tool_calling_agent
            prompt = ChatPromptTemplate.from_messages([
                ("system", self.get_system_prompt()),
                ("human", "{input}"),
                MessagesPlaceholder(variable_name="agent_scratchpad")
            ])
            
            self.agent_executor = create_tool_calling_agent(self.llm, self.tools, prompt)
        except Exception as e:
            logger.warning(f"create_tool_calling_agent failed: {e}, using fallback")
            # 使用基础的AgentExecutor作为备选方案
            from langchain.agents import AgentExecutor
            from langchain.chains import LLMChain
            
            # 创建一个简单的LLM链
            llm_chain = LLMChain(llm=self.llm, prompt=ChatPromptTemplate.from_messages([
                ("system", self.get_system_prompt()),
                ("human", "{input}")
            ]))
            
            self.agent_executor = AgentExecutor.from_agent_and_tools(
                agent=llm_chain,
                tools=self.tools,
                verbose=True
            )
        
        logger.info(f"数据库智能体 {self.name} 初始化完成")

class FinanceAgent(Agent):
    """金融数据库智能体"""
    
    def __init__(self, name: str = "Finance Agent", description: str = "可以访问和查询金融数据库的智能体"):
        super().__init__(name, description)
        self.finance_db_tool = FinanceDatabaseTool()
        self.schema_reader = DatabaseSchemaReader()
        self._setup_tools()
        self.initialize_agent()
    
    def _setup_tools(self):
        """设置工具"""
        # 获取表名工具
        def get_finance_table_names(*args, **kwargs) -> List[str]:
            """获取金融数据库中所有表名"""
            # 忽略传入的任何参数，保持向后兼容
            return self.finance_db_tool.get_table_names()
        
        get_finance_table_names_tool = Tool.from_function(
            func=get_finance_table_names,
            name="get_finance_table_names",
            description="获取金融数据库中所有表名"
        )
        
        # 获取表结构工具
        def get_finance_table_schema(table_name: str, *args, **kwargs) -> List[Dict[str, Any]]:
            """获取金融表的结构信息"""
            return self.finance_db_tool.get_table_schema(table_name)
        
        get_finance_table_schema_tool = Tool.from_function(
            func=get_finance_table_schema,
            name="get_finance_table_schema",
            description="获取金融表的结构信息。参数：table_name(表名)"
        )
        
        # 获取表示例数据工具
        def get_finance_table_sample_data(table_name: str, limit: int = 5, *args, **kwargs) -> List[Dict[str, Any]]:
            """获取金融表的示例数据"""
            return self.finance_db_tool.get_table_sample_data(table_name, limit)
        
        get_finance_table_sample_data_tool = Tool.from_function(
            func=get_finance_table_sample_data,
            name="get_finance_table_sample_data",
            description="获取金融表的示例数据。参数：table_name(表名), limit(返回记录数，默认5条)"
        )
        
        # 执行查询工具
        def execute_finance_query(query: str, *args, **kwargs) -> List[Dict[str, Any]]:
            """执行金融数据库查询"""
            return self.finance_db_tool.execute_query(query)
        
        execute_finance_query_tool = Tool.from_function(
            func=execute_finance_query,
            name="execute_finance_query",
            description="执行金融数据库查询。参数：query(SQL查询语句)"
        )
        
        # 读取数据库结构文档工具
        def read_database_schema() -> str:
            """读取数据库结构文档"""
            return self.schema_reader.read_schema_document()
        
        read_database_schema_tool = Tool.from_function(
            func=read_database_schema,
            name="read_database_schema",
            description="读取完整的数据库结构文档，包含所有表的详细信息"
        )
        
        # 获取特定表信息工具
        def get_table_info(table_name: str) -> str:
            """获取特定表的详细信息"""
            return self.schema_reader.get_table_info(table_name)
        
        get_table_info_tool = Tool.from_function(
            func=get_table_info,
            name="get_table_info",
            description="获取特定表的详细信息，包括表结构和示例数据。参数：table_name(表名)"
        )
        
        self.tools = [
            get_finance_table_names_tool,
            get_finance_table_schema_tool,
            get_finance_table_sample_data_tool,
            execute_finance_query_tool,
            read_database_schema_tool,
            get_table_info_tool
        ]
        
    def get_system_prompt(self) -> str:
        """获取金融数据库智能体的系统提示词"""
        return """
你是一个专业的金融数据专家助手，可以精准访问和查询金融数据库。数据库包含以下10张表，数据区间为2019年至2021年：

1. 基金基本信息 - 包含基金名称、代码、成立日期、管理人、托管人、费率等基本信息
2. 基金股票持仓明细 - 包含基金持有的股票信息、持仓数量、市值、占比等
3. 基金债券持仓明细 - 包含基金持有的债券信息、持仓数量、市值、占比等
4. 基金可转债持仓明细 - 包含基金持有的可转债信息、持仓数量、市值、占比等
5. 基金日行情表 - 包含基金每日的净值、涨跌幅、成交量等数据
6. A股票日行情表 - 包含A股每日的开盘价、收盘价、最高价、最低价、成交量、成交额等数据
7. 港股票日行情表 - 包含港股每日的开盘价、收盘价、最高价、最低价、成交量、成交额等数据
8. A股公司行业划分表 - 包含A股公司所属的行业分类信息
9. 基金规模变动表 - 包含基金规模的变动情况
10. 基金份额持有人结构 - 包含基金份额持有人的结构信息

关于数据库的详细结构信息，请参考文档：llm_backend/docs/database_schema.md。该文档包含了每张表的详细字段说明、数据类型以及示例数据，可以作为构建SQL查询的重要参考。

你可以使用以下工具：

1. get_finance_table_names - 获取金融数据库中所有表名，用于了解有哪些可用的数据表
2. get_finance_table_schema - 获取金融表的结构信息，包括列名、数据类型等，帮助构建准确的SQL查询
3. get_finance_table_sample_data - 获取金融表的示例数据，了解数据格式和内容
4. execute_finance_query - 执行金融数据库查询，需要提供正确的SQL查询语句
5. read_database_schema - 读取完整的数据库结构文档
6. get_table_info - 获取特定表的详细信息，包括表结构和示例数据

使用指南（请严格遵循）：
1. 在构建复杂查询前，参考 llm_backend/docs/database_schema.md 文档了解表结构和关联关系
2. 对于任何查询，第一步必须使用get_finance_table_names获取所有表名，确保了解可用的数据资源
3. 第二步使用get_finance_table_schema查看与查询相关的表的结构，特别关注列名和数据类型
4. 第三步构建SQL查询语句时，必须严格按照表结构中的列名，使用双引号包裹包含特殊字符的列名（如收盘价(元)）
5. 执行查询时，确保SQL语法正确，尤其是WHERE条件中的日期格式要与表中数据一致
6. 查询结果应限制在最多10条记录，对于聚合查询可以返回所有结果

SQL查询技巧：
1. 所有SQL查询都必须是SELECT语句，禁止执行INSERT、UPDATE、DELETE等修改数据的操作
2. 对于数值比较，确保使用正确的比较运算符（=, >, <, >=, <=）
3. 对于日期筛选，确保日期格式与表中数据一致（通常为YYYYMMDD格式）
4. 使用ORDER BY子句对结果进行排序，使用LIMIT限制返回的记录数
5. 对于统计查询，使用SUM、AVG、MAX、MIN等聚合函数
6. 对于多表查询，确保正确使用JOIN语句和连接条件，参考文档中的表关联关系部分

回答规范：
1. 回答时必须基于查询结果，确保准确性和客观性
2. 回答要简洁明了，重点突出，避免冗长的解释
3. 如果查询没有返回结果，请明确说明未找到相关数据
4. 如果查询出错，请分析错误原因并提供可能的解决方案
5. 对于复杂的多步骤查询，必须详细展示每一步的结果和最终结论

现在开始处理用户的查询，请严格按照上述要求执行操作。"""

    def initialize_agent(self):
        """初始化金融数据库智能体"""
        try:
            # 提前导入AgentExecutor以避免局部变量引用错误
            from langchain.agents import AgentExecutor
            
            # 创建提示模板
            prompt = ChatPromptTemplate.from_messages([
                ("system", self.get_system_prompt()),
                ("human", "{input}"),
                MessagesPlaceholder(variable_name="agent_scratchpad")
            ])
            
            # 使用create_tool_calling_agent创建代理并包装到AgentExecutor中
            agent = create_tool_calling_agent(self.llm, self.tools, prompt)
            self.agent_executor = AgentExecutor.from_agent_and_tools(
                agent=agent,
                tools=self.tools,
                verbose=True,
                return_intermediate_steps=False,  # 禁用中间步骤返回，避免'tuple'错误
                handle_parsing_errors=True         # 启用解析错误处理
            )
        except Exception as e:
            logger.warning(f"create_tool_calling_agent failed: {e}, using fallback")
            # 备选方案：使用基础的AgentExecutor
            from langchain.agents import AgentExecutor
            
            # 创建一个详细的提示模板
            prompt = ChatPromptTemplate.from_messages([
                ("system", self.get_system_prompt()),
                ("human", "{input}"),
                MessagesPlaceholder(variable_name="agent_scratchpad", optional=True)
            ])
            
            # 创建LLM链 (使用新的语法替代已弃用的LLMChain)
            llm_chain = prompt | self.llm
            
            # 创建AgentExecutor
            self.agent_executor = AgentExecutor.from_agent_and_tools(
                agent=llm_chain,
                tools=self.tools,
                verbose=True,
                return_intermediate_steps=False,  # 禁用中间步骤返回
                handle_parsing_errors=True         # 启用解析错误处理
            )
        
        logger.info(f"金融数据库智能体 {self.name} 初始化完成")

# MCP兼容功能
class MCPDatabaseAgent:
    """MCP兼容的数据库智能体"""
    
    def __init__(self, db_path: str = None):
        self.db_tool = DatabaseTool(db_path)
        self.llm = get_llm()
    
    def list_databases(self) -> List[str]:
        """列出所有数据库"""
        return mcp_list_databases(self.db_tool)
    
    def execute_query(self, db_name: str, query: str) -> List[Dict[str, Any]]:
        """执行查询"""
        return mcp_execute_query(self.db_tool, db_name, query)
    
    def get_table_names(self, db_name: str) -> List[str]:
        """获取表名"""
        return mcp_get_table_names(self.db_tool, db_name)
    
    def get_table_schema(self, db_name: str, table_name: str) -> List[Dict[str, Any]]:
        """获取表结构"""
        return mcp_get_table_schema(self.db_tool, db_name, table_name)